import os
import glob
import json

import cv2
import numpy as np
from tqdm import tqdm


def main():
    crop_dir = r'D:\lliujian\other\新建文件夹\测试照片_crop'
    stitch_dir = os.path.join(os.path.dirname(crop_dir), os.path.basename(crop_dir) + '_stitch')

    img_suffix = '.bmp'
    label = 'Needle'

    if not os.path.exists(stitch_dir):
        os.makedirs(stitch_dir)
    for i, pos_pth in enumerate(tqdm(glob.glob(crop_dir + f'\\*_positions.json'))):
        with open(pos_pth, 'r', encoding='utf-8') as f:
            pos_data = json.load(f)
        name = pos_data['name'].replace(img_suffix, '')
        crop_shape = pos_data['crop_shape']
        img_paths = glob.glob(crop_dir + f'\\{name}*{img_suffix}')
        stitch_mask = np.zeros(pos_data['image_shape'][:2], np.uint8)

        for j, pos in enumerate(pos_data['positions']):
            crop_json_path = os.path.join(crop_dir, name.replace(img_suffix, '') + f'_{j}.json')
            with open(crop_json_path, 'r', encoding='utf-8') as f:
                shapes = json.load(f)['shapes']
            polygons = []
            for shape in shapes:
                if shape['shape_type'] == 'polygon' and shape['label'] == label:
                    poly = shape['points']
                    polygons.append(np.array(poly, dtype=np.int32))
                else:
                    print("Skip mismatched label: shape_type={}  label={}".format(shape['shape_type'],
                                                                                  shape['label']))
            crop_msk = np.zeros(crop_shape[:2], np.uint8)
            cv2.fillPoly(crop_msk, polygons, [255])
            l, t, r, b = pos
            stitch_mask[t:b, l:r] = crop_msk
            # cv2.imencode(img_suffix, crop_msk)[1].tofile(os.path.join(stitch_dir, '{}_{}'.format(j, pos_data['name'])))

        cv2.imencode(img_suffix, stitch_mask)[1].tofile(os.path.join(stitch_dir, pos_data['name']))
        # break


if __name__ == "__main__":
    main()